import matplotlib.pyplot as plt
import numpy as np

import matplotlib

matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文

matplotlib.rcParams['axes.unicode_minus'] = False  # 正常显示负号

def plot_scatter_with_labels(names, values, colors=None, title="Scatter Plot", figsize=(10, 6), 
                            save_path=None, connect_points=False, line_style='-', line_width=1, 
                            line_color='gray', line_alpha=0.7):
    """
    绘制带标签的散点图
    
    参数:
    names: list of str, 点的名称列表
    values: list of float/int, 点的值列表
    colors: list of str/tuple, 可选，点的颜色列表，可以是颜色名称或RGB元组
    title: str, 图表标题
    figsize: tuple, 图表大小
    save_path: str, 可选，保存图片的路径
    connect_points: bool, 是否连接散点，默认为False
    line_style: str, 连接线的样式，默认为'-'，可选：'-', '--', '-.', ':', 等
    line_width: float, 连接线的宽度，默认为1
    line_color: str/tuple, 连接线的颜色，默认为'gray'
    line_alpha: float, 连接线的透明度，默认为0.7
    """
    # 创建图表
    plt.figure(figsize=figsize)
    
    # 生成x坐标（均匀分布）
    x = np.arange(len(values))
    
    # 如果需要连接点，先绘制线条
    if connect_points and len(values) > 1:
        plt.plot(x, values, linestyle=line_style, linewidth=line_width, 
                color=line_color, alpha=line_alpha, zorder=1)
    
    # 处理颜色参数
    if colors is None:
        # 默认颜色
        scatter = plt.scatter(x, values, alpha=0.6, zorder=2)
    else:
        # 确保颜色列表长度与数据点数量一致
        if len(colors) < len(values):
            colors = colors * (len(values) // len(colors) + 1)
            colors = colors[:len(values)]
        # 使用指定颜色
        scatter = plt.scatter(x, values, c=colors, alpha=0.6, zorder=2)
    
    # 添加标签
    for i, (name, value) in enumerate(zip(names, values)):
        plt.annotate(name, (x[i], value), 
                    xytext=(5, 5), 
                    textcoords='offset points',
                    fontsize=8)
    
    # 设置标题和轴标签
    plt.title(title)
    plt.xlabel("Index")
    plt.ylabel("Value")
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图片（如果指定了保存路径）
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"图表已保存至: {save_path}")
    
    # 显示图表
    plt.show()


def plot_multi_scatter(data_groups, title="Multi-Series Scatter Plot", figsize=(12, 8), 
                      save_path=None, connect_points=False, line_styles=None, 
                      marker_sizes=None, alpha=0.7, add_legend=True, x_label="Index", y_label="Value"):
    """
    在同一张图上绘制多组散点数据
    
    参数:
    data_groups: list of dict, 每个字典包含一组数据，格式如下:
                {
                    'names': list of str,  # 点的名称列表
                    'values': list of float/int,  # 点的值列表
                    'colors': list of str/tuple,  # 可选，点的颜色列表
                    'label': str,  # 可选，该组数据的标签（用于图例）
                    'line_color': str/tuple,  # 可选，连接线的颜色
                    'line_style': str,  # 可选，连接线的样式
                    'line_width': float,  # 可选，连接线的宽度
                }
    title: str, 图表标题
    figsize: tuple, 图表大小
    save_path: str, 可选，保存图片的路径
    connect_points: bool, 是否连接散点，默认为False（可被每组数据中的设置覆盖）
    line_styles: list of str, 可选，连接线样式列表，如果未在每组数据中指定，则循环使用此列表
    marker_sizes: list of float, 可选，点大小列表，如果未在每组数据中指定，则循环使用此列表
    alpha: float, 点和线的透明度，默认为0.7
    add_legend: bool, 是否添加图例，默认为True
    x_label: str, x轴标签
    y_label: str, y轴标签
    """
    # 创建图表
    plt.figure(figsize=figsize)
    
    # 默认线条样式
    if line_styles is None:
        line_styles = ['-', '--', '-.', ':']
    
    # 默认点大小
    if marker_sizes is None:
        marker_sizes = [30, 40, 50, 60, 70]
    
    # 默认颜色循环
    default_colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
    
    # 图例项
    legend_elements = []
    
    # 绘制每组数据
    for i, group in enumerate(data_groups):
        # 提取数据
        names = group.get('names', [])
        values = group.get('values', [])
        
        if not names or not values:
            continue
            
        # 确保names和values长度一致
        min_len = min(len(names), len(values))
        names = names[:min_len]
        values = values[:min_len]
        
        # 生成x坐标（均匀分布）
        x = np.arange(len(values))
        
        # 获取组标签
        group_label = group.get('label', f'Series {i+1}')
        
        # 获取颜色
        colors = group.get('colors')
        if colors is None:
            # 使用默认颜色
            color = default_colors[i % len(default_colors)]
            colors = [color] * len(values)
        elif isinstance(colors, (str, tuple)):
            # 如果colors是单个颜色，转换为列表
            colors = [colors] * len(values)
        
        # 确保颜色列表长度与数据点数量一致
        if len(colors) < len(values):
            colors = colors * (len(values) // len(colors) + 1)
            colors = colors[:len(values)]
        
        # 获取点大小
        marker_size = group.get('marker_size', marker_sizes[i % len(marker_sizes)])
        
        # 是否连接点
        group_connect = group.get('connect_points', connect_points)
        
        # 如果需要连接点，先绘制线条
        if group_connect and len(values) > 1:
            line_color = group.get('line_color', colors[0] if isinstance(colors[0], str) else default_colors[i % len(default_colors)])
            line_style = group.get('line_style', line_styles[i % len(line_styles)])
            line_width = group.get('line_width', 1.5)
            
            plt.plot(x, values, linestyle=line_style, linewidth=line_width, 
                    color=line_color, alpha=alpha, zorder=1, label=group_label)
            
            # 已经添加了标签，不需要在散点图中再添加
            scatter_label = None
        else:
            # 如果没有连接点，则在散点图中添加标签
            scatter_label = group_label
        
        # 绘制散点图
        scatter = plt.scatter(x, values, c=colors, s=marker_size, alpha=alpha, 
                             zorder=2, label=scatter_label)
        
        # 添加点标签
        for j, (name, value) in enumerate(zip(names, values)):
            if name:  # 只有当名称不为空时才添加标签
                plt.annotate(name, (x[j], value), 
                            xytext=(5, 5), 
                            textcoords='offset points',
                            fontsize=8)
    
    # 添加图例
    if add_legend:
        plt.legend()
    
    # 设置标题和轴标签
    plt.title(title)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图片（如果指定了保存路径）
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"图表已保存至: {save_path}")
    
    # 显示图表
    plt.show()


# 使用示例
if __name__ == "__main__":
    # 单组数据示例
    test_names = ["A", "B", "C", "D", "E"]
    test_values = [10, 25, 15, 30, 20]
    test_colors = ['red', 'green', 'blue', 'purple', 'orange']
    
    # # 绘制单组数据示例
    # plot_scatter_with_labels(test_names, test_values, colors=test_colors, title="散点图示例")
    
    # # 绘制连接线示例
    # plot_scatter_with_labels(test_names, test_values, colors=test_colors, 
    #                        title="连接线散点图示例", connect_points=True)
    
    # 多组数据示例
    group1 = {
        'names': ['A1', 'B1', 'C1', 'D1', 'E1'],
        'values': [10, 25, 15, 30, 20],
        'colors': 'red',
        'label': '组1',
        'connect_points': True,
        'line_style': '-'
    }
    
    group2 = {
        'names': ['A2', 'B2', 'C2', 'D2', 'E2'],
        'values': [15, 30, 25, 10, 35],
        'colors': 'blue',
        'label': '组2',
        'connect_points': True,
        'line_style': '--'
    }
    
    group3 = {
        'names': ['第一', '语文', '数学', '英语'],
        'values': [5, 20, 10, 25],
        'colors': 'green',
        'label': '组3',
        'connect_points': False
    }
    
    # 绘制多组数据示例
    plot_multi_scatter([group1, group2, group3], title="多组数据散点图示例")
